﻿using System;
using System.Collections.Generic;
using System.IO;

namespace tuimg
{
    using tuimg.element;
    using tuimg.helpers;
    using tuimg.structs;

    public class TuimgPng : IDisposable
    {
        internal static readonly ulong png_tag = 0x0A1A0A0D474E5089;
        internal static readonly uint ihdr_tag = 0x52444849;
        internal struct_png_IHDR ihdr;
        internal struct_png_PLTE plte;
        internal struct_png_IEND iend;
        internal struct_png_sRGB srgb;
        internal struct_png_gAMA gama;
        internal struct_png_pHYs phys;
        internal Stream imgdata;
        internal ushort bitcount;
        internal uint onerowbytes;
        public TuimgPng(string pngpath)
        {
            imgdata = new MemoryStream();
            using (var fr = File.Open(pngpath, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
            {
                if (png_tag != helper.Readulong(fr)) throw new ArgumentException("not png file.");
                load_all_block(fr);

                calc_imgdata_info();
                imgdata_to_color32_or_24();
                imgdata_filters_tozero();
            }
        }
        public TuimgPng(Stream s)
        {
            imgdata = new MemoryStream();

            if (png_tag != helper.Readulong(s)) throw new ArgumentException("not png file.");
            load_all_block(s);

            calc_imgdata_info();
            imgdata_to_color32_or_24();
            imgdata_filters_tozero();
        }
        public TuimgPng(int width, int height, ushort bitcount = 32)
        {
            imgdata = new MemoryStream();
            ihdr = new struct_png_IHDR { };
            ihdr.length = 13;
            ihdr.tag = "IHDR";
            ihdr.width = (uint)width;
            ihdr.height = (uint)height;
            ihdr.color_depth = 8;
            ihdr.color_type = 6;
            ihdr.png_spec_compression_method= 0;
            ihdr.png_spec2_filter_method = 0;
            ihdr.not_interlaced = 0;

            srgb = new struct_png_sRGB { };
            srgb.length = 1;
            srgb.tag = "sRGB";
            srgb.rendering_intent = 0;

            gama = new struct_png_gAMA { };
            gama.length = 4;
            gama.tag = "gAMA";
            gama.data = 2410741760;

            phys = new struct_png_pHYs { };
            phys.length = 9;
            phys.tag = "pHYs";
            phys.unit_specifier = 1;
            phys.x_axis = 0;
            phys.y_axis = 0;

            this.bitcount = bitcount;
            calc_imagedata_onerow();
            var buf = new byte[onerowbytes];
            for (var i = 1; i < onerowbytes; i++) buf[i] = 0;
            for (var i = 0; i < height; i++) imgdata.Write(buf, 0, (int)onerowbytes);
        }
        public long Height { get { return ihdr.height; } }
        public long Width { get { return ihdr.width; } }
        public Stream ToStream() {
            var mem = new MemoryStream();
            ToStream(mem);
            return mem;
        }
        public void ToStream(Stream mem)
        {
            long ix = 0;
            helper.Write(mem, png_tag);
            helper.Write2(mem, ihdr.length);
            ix = mem.Position;
            helper.Writechars(mem, ihdr.tag);
            helper.Write2(mem, ihdr.width);
            helper.Write2(mem, ihdr.height);
            helper.Write(mem, ihdr.color_depth);
            helper.Write(mem, ihdr.color_type);
            helper.Write(mem, ihdr.png_spec_compression_method);
            helper.Write(mem, ihdr.png_spec2_filter_method);
            helper.Write(mem, ihdr.not_interlaced);
            helper.write_crc(mem, ix, ihdr.length + 4);

            if (srgb.length > 0)
            {
                helper.Write2(mem, srgb.length);
                ix = mem.Position;
                helper.Writechars(mem, srgb.tag);
                helper.Write(mem, srgb.rendering_intent);
                helper.write_crc(mem, ix, srgb.length + 4);
            }

            if (gama.length > 0)
            {
                helper.Write2(mem, gama.length);
                ix = mem.Position;
                helper.Writechars(mem, gama.tag);
                helper.Write(mem, gama.data);
                helper.write_crc(mem, ix, gama.length + 4);
            }

            if (phys.length > 0)
            {
                helper.Write2(mem, phys.length);
                ix = mem.Position;
                helper.Writechars(mem, phys.tag);
                helper.Write(mem, phys.unit_specifier);
                helper.Write(mem, phys.x_axis);
                helper.Write(mem, phys.y_axis);
                helper.write_crc(mem, ix, phys.length + 4);
            }

            // 压缩
            filter_imgdata();
            using (var compress_imgdata = new MemoryStream())
            {
                helper.defalte_compress(imgdata, compress_imgdata);

                var first_compressbag = true; // 首数据包加长
                var len = compress_imgdata.Length >> 16;
                for (; len>0; len--)
                {
                    if (first_compressbag)
                    {
                        helper.Write2(mem, (uint)65537);
                        ix = mem.Position;
                        helper.Writechars(mem, "IDAT");
                        helper.Write(mem, struct_png_IDAT.ZLIB_Deflate_32K);
                        helper.Write(mem, struct_png_IDAT.ZLIB_FLEVEL_2);
                        helper.ReadToMemory(compress_imgdata, mem, 65535);
                        helper.write_crc(mem, ix, 65541); // 65537+4
                        first_compressbag = false;
                    }
                    else
                    {
                        helper.Write2(mem, (uint)65535);
                        ix = mem.Position;
                        helper.Writechars(mem, "IDAT");
                        helper.ReadToMemory(compress_imgdata, mem, 65535);
                        helper.write_crc(mem, ix, 65539);//65535+4
                    }
                }

                len = compress_imgdata.Length & 0xffff;
                if (len > 0)
                {
                    if (first_compressbag)
                    {
                        helper.Write2(mem, (uint)len+2);
                        ix = mem.Position;
                        helper.Writechars(mem, "IDAT");
                        helper.Write(mem, struct_png_IDAT.ZLIB_Deflate_32K);
                        helper.Write(mem, struct_png_IDAT.ZLIB_FLEVEL_2);
                        helper.ReadToMemory(compress_imgdata, mem, (uint)len);
                        helper.write_crc(mem, ix, (uint)(len + 6));
                    }
                    else
                    {

                        helper.Write2(mem, (uint)len);
                        ix = mem.Position;
                        helper.Writechars(mem, "IDAT");
                        helper.ReadToMemory(compress_imgdata, mem, (uint)len);
                        helper.write_crc(mem, ix, (uint)(len + 4));
                    }
                }
            }


            imgdata_filters_tozero();

            helper.Write2(mem, 0);
            ix = mem.Position;
            helper.Writechars(mem, "IEND");
            helper.write_crc(mem, ix, 4);
            mem.Flush();
            mem.Seek(0, SeekOrigin.Begin);
        }
        public void DrawGrid(TuimgGrid grid, TuimgColor color)
        {
            DrawGrid(0, 0, grid, color);
        }
        public void DrawGrid(int x, int y, TuimgGrid grid)
        {
            DrawGrid(x, y, grid, TuimgColor.Black);
        }
        public void DrawGrid(int x, int y, TuimgGrid grid, TuimgColor color)
        {
            var _color = color.ToBitColor(bitcount);
            int i, j, _y, _x;
            for (i = 0; i < grid.Height; i++) for (j = 0, _y = i + y, _x = x; j < grid.Width; j++, _x++) if (grid.Read(j, i)) draw_point(_x, _y, _color);
        }
        public void DrawSolidBox(int x, int y, int width, int height, TuimgColor color)
        {
            var _color = color.ToBitColor(bitcount);
            for (var i = 0; i < height && (y + i) < height; i++)
            {
                brush_goto(x, y + i);
                for (var j = 0; j < width && (j + x) < width; j++) write_color(_color);
            }
        }
        internal void draw_point(int x, int y, object color)
        {
            brush_goto(x, y);
            helper.WriteColorRgb(imgdata, color);
        }
        internal void calc_imgdata_info()
        {
            bitcount = (ushort)(((imgdata.Length / ihdr.height - 1) << 3) / ihdr.width);
            calc_imagedata_onerow();
        }
        public TuimgBmp ToBmp(TuimgColor background)
        {
            var bmp = new TuimgBmp((int)ihdr.width, (int) ihdr.height);
            for (var i = 0; i < ihdr.height; i++)
            {
                bmp.brush_goto(0, i);
                brush_goto(0, i);
                for (var j = 0; j < ihdr.width; j++) bmp.WriteColor(helper.ReadColor(imgdata, (short)bitcount, background));
            }
            return bmp;
        }
        public TuimgGrid ToAndGrid(int transport =0)
        {
            var grid = new TuimgGrid((int)ihdr.width, (int)ihdr.height);

            for (var i = 0; i < ihdr.height; i++)
            {
                brush_goto(0, i);
                for (var j = 0; j < ihdr.width; j++) if (!helper.IsTransparent(imgdata, (short)bitcount, transport)) grid.On(j, i);
            }
            return grid;
        }
        internal void imgdata_to_color32_or_24()
        {
            // [逐步更新] 处理颜色索引的情况
        }
        internal void filter_imgdata() //滤波数据
        {
            byte filter_mode = 4;
            var _buf = new byte[onerowbytes];
            var buf = new byte[onerowbytes];
            var buf2 = new byte[onerowbytes];
            var color_bytecount = (bitcount >> 3);
            for (var i = 0; i < ihdr.height; i++)
            {
                goto_row_filtermode(i);
                if (filter_mode == 0)
                {
                    imgdata.Read(_buf, 0, (int)(onerowbytes));
                    continue;
                }
                imgdata.Read(buf2, 0, (int)(onerowbytes));

                // [逐步更新]滤波策略 智能选择最小压缩
                buf[0] = filter_mode;

                var j = color_bytecount + 1;
                switch (filter_mode)
                {
                    case 1: // x-a
                        for (var jj = 1; jj < j; jj++) buf[jj] = buf2[jj];
                        for (; j < onerowbytes; j++) buf[j] = (byte)(buf2[j] - buf2[j - color_bytecount]);

                        Array.Copy(buf2, _buf, onerowbytes);
                        goto_row_filtermode(i);
                        imgdata.Write(buf, 0, (int)onerowbytes);
                        break;
                    case 2:
                        for (var jj = 1; jj < j; jj++) buf[jj] = (byte)(buf2[jj]-_buf[jj]);
                        for (; j < onerowbytes; j++) buf[j] = (byte)(buf2[j] - _buf[j]);

                        Array.Copy(buf2, _buf, onerowbytes);
                        goto_row_filtermode(i);
                        imgdata.Write(buf, 0, (int)onerowbytes);
                        break;
                    case 3:
                        for (var jj = 1; jj < j; jj++) buf[jj] = (byte)(buf2[jj]-helper.round_down(_buf[jj]/2));
                        for (; j < onerowbytes; j++) buf[j] = (byte)(buf2[j] - helper.round_down((_buf[j] + buf2[j-color_bytecount]) / 2));

                        Array.Copy(buf2, _buf, onerowbytes);
                        goto_row_filtermode(i);
                        imgdata.Write(buf, 0, (int)onerowbytes);
                        break;
                    case 4:
                        for (var jj = 1; jj < j; jj++) buf[jj] = (byte)(buf2[jj]- helper.paethPrediction(0, _buf[jj], 0));
                        for (; j < onerowbytes; j++) buf[j] = (byte)(buf2[j] - helper.paethPrediction(buf2[j - color_bytecount], _buf[j], _buf[j - color_bytecount]));

                        Array.Copy(buf2, _buf, onerowbytes);
                        goto_row_filtermode(i);
                        imgdata.Write(buf, 0, (int)onerowbytes);
                        break;
                    default:
                        break;
                }
            }
        }
        internal void imgdata_filters_tozero()
        {
            byte filter_mode;
            int j;
            var buf = new byte[onerowbytes];
            var buf2 = new byte[onerowbytes];
            var color_bytecount = (bitcount>>3);

            for (var i =0; i < ihdr.height; i++)
            {
                goto_row_filtermode(i);
                filter_mode = helper.Readbyte(imgdata);
                if (filter_mode == 0)
                {
                    buf[0] = filter_mode;
                    imgdata.Read(buf, 1, (int)(onerowbytes - 1));
                    continue;
                }

                imgdata.Read(buf2, 1, (int)(onerowbytes-1));
                buf2[0] = 0;
                j = color_bytecount + 1;
                switch (filter_mode)
                {
                    case 1: // x-a
                        for (var jj = 1; jj < j; jj++) buf2[jj] = (byte)(helper.toint8(buf2[jj]));
                        for (; j < onerowbytes; j++) buf2[j] = (byte)(helper.toint8(buf2[j]) + buf2[j - color_bytecount]);
                        Array.Copy(buf2, buf, onerowbytes);

                        goto_row_filtermode(i);
                        imgdata.Write(buf2, 0, (int)onerowbytes);
                        break;
                    case 2:
                        for (var jj = 1; jj < j; jj++) buf2[jj] = (byte)(buf[jj] + helper.toint8(buf2[jj]));
                        for (; j < onerowbytes; j++) buf2[j] = (byte)(helper.toint8(buf2[j]) +buf[j]);
                        Array.Copy(buf2, buf, onerowbytes);

                        goto_row_filtermode(i);
                        imgdata.Write(buf2, 0, (int)onerowbytes);
                        break;
                    case 3:
                        for (var jj = 1; jj < j; jj++) buf2[jj] = (byte)(helper.round_down((buf[jj]) / 2) + helper.toint8(buf2[jj]));
                        for (; j < onerowbytes; j++) buf2[j] = (byte)(helper.toint8(buf2[j]) + helper.round_down((buf[j] + buf2[j - color_bytecount]) / 2));
                        Array.Copy(buf2, buf, onerowbytes);

                        goto_row_filtermode(i);
                        imgdata.Write(buf2, 0, (int)onerowbytes);
                        break;
                    case 4:
                        for (var jj = 1; jj < j; jj++) buf2[jj] = (byte)(helper.paethPrediction(0, buf[jj],0) + helper.toint8(buf2[jj]));
                        for (; j < onerowbytes; j++) buf2[j] = (byte)(helper.toint8(buf2[j]) + helper.paethPrediction(buf2[j - color_bytecount], buf[j], buf[j - color_bytecount]));
                        Array.Copy(buf2, buf, onerowbytes);

                        goto_row_filtermode(i);
                        imgdata.Write(buf2, 0, (int)onerowbytes);
                        break;
                    default:
                        break;
                }
            }
        }
        internal void goto_row_filtermode(int row)
        {
            imgdata.Position = row* onerowbytes;
        }
        internal void calc_imagedata_onerow()
        {
            onerowbytes = (bitcount * ihdr.width +15) >> 3;
        }
        internal void brush_goto(int x, int y)
        {
            imgdata.Position = onerowbytes*y +((x*bitcount +15)>>3);
        }
        internal void brush_goto_next()
        {
            imgdata.Position += bitcount>>3;
        }
        public void WriteColor(TuimgColor color)
        {
            write_color(color.ToBitColor(bitcount));
        }
        public TuimgColor ReadColor()
        {
            return helper.ReadColor(imgdata, (short)bitcount);
        }
        internal void write_color(object color)
        {
            helper.WriteColorRgb(imgdata, color);
        }
        internal void load_all_block(Stream fr)
        {
            using (var _mem = new MemoryStream())
            {
                foreach (var block in read_blocks(fr))
                {
                    switch (block.tag)
                    {
                        case "IHDR":
                            ihdr = read_ihdr(fr, block);
                            break;
                        case "PLTE":
                            plte = read_plte(fr, block);
                            break;
                        case "sRGB":
                            srgb = read_sRGB(fr, block);
                            break;
                        case "gAMA": // 图像y数据
                            gama = read_gAMA(fr, block);
                            break;
                        case "pHYs": // 物理像素尺寸
                            phys = read_pHYs(fr, block);
                            break;
                        case "cHRM": // 基色和白色点
                            break;
                        case "sBIT": // 样本有效位数据
                            break;
                        case "bKGD": // 背景数据
                            break;
                        case "hIST": // 直方图数据
                            break;
                        case "tRNS": // 透明数据
                            break;
                        case "oFFs": // 公共
                            break;
                        case "sCAL": // 公共
                            break;
                        case "tIME": // 最后修改时间
                            break;
                        case "tEXt": // 文本信息
                            break;
                        case "zTXt": // 压缩文本
                            break;
                        case "fRAc": // 公共
                            break;
                        case "gIFg": // 公共
                            break;
                        case "gIFt": // 公共
                            break;
                        case "gIFx": // 公共
                            break;
                        case "IDAT": // 图像数据
                            var idat = read_IDAT(fr, block, _mem);
                            break;
                        case "IEND": // 图像结束
                            iend = read_iend(fr, block);
                            helper.deflate_decompress(_mem, imgdata);
                            imgdata.Seek(0, SeekOrigin.Begin);
                            break;
                    }
                }
            }
        }
        internal IEnumerable<struct_png_datablock> read_blocks(Stream sr)
        {
            while (sr.Position < sr.Length)
            {
                var block = new struct_png_datablock { };
                block.length = helper.Readuint2(sr);
                block.pos = sr.Position;
                block.crc = helper.calc_crc(sr, block.pos, block.length + 4);
                sr.Position = block.pos;
                block.tag = helper.Readchars(sr, 4);
                yield return block;
                if (block.tag == "IEND")
                    break;
                sr.Position = block.pos+ block.length + 8; // tag 4-byte + crc 4-byte
            }
        }
        internal struct_png_IHDR read_ihdr(Stream sr, struct_png_datablock b)
        {
            var ihdr = new struct_png_IHDR { };
            ihdr.length = b.length;
            ihdr.tag = b.tag;
            ihdr.width = helper.Readuint2(sr);
            ihdr.height = helper.Readuint2(sr);
            ihdr.color_depth = helper.Readbyte(sr);
            ihdr.color_type = helper.Readbyte(sr);
            ihdr.png_spec_compression_method = helper.Readbyte(sr);
            ihdr.png_spec2_filter_method = helper.Readbyte(sr);
            ihdr.not_interlaced = helper.Readbyte(sr);
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png ihdr error crc.");
            return ihdr;
        }
        internal struct_png_PLTE read_plte(Stream sr, struct_png_datablock b)
        {
            var plte = new struct_png_PLTE { };
            plte.length = b.length;
            plte.tag = b.tag;
            if ((b.length & 0b11) != 0) throw new ArgumentException("plte data error.");
            plte.color_length = (byte)(b.length / 3);
            plte.colors = new TuimgColor[plte.color_length];
            for (int i = 0; i < plte.colors.Length; i++) plte.colors[i] = helper.ReadColor24(sr);
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png plte error crc.");
            return plte;
        }
        internal struct_png_IEND read_iend(Stream sr, struct_png_datablock b)
        {
            var iend = new struct_png_IEND { };

            iend.length = b.length;
            iend.tag = b.tag;
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png iend error crc.");
            return iend;
        }
        internal struct_png_sRGB read_sRGB(Stream sr, struct_png_datablock b)
        {
            var sRGB = new struct_png_sRGB { };

            sRGB.length = b.length;
            sRGB.tag = b.tag;
            sRGB.rendering_intent = helper.Readbyte(sr);
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png sRGB error crc.");
            return sRGB;
        }
        internal struct_png_gAMA read_gAMA(Stream sr, struct_png_datablock b)
        {
            var gAMA = new struct_png_gAMA { };

            gAMA.length = b.length;
            gAMA.tag = b.tag;
            gAMA.data = helper.Readuint(sr);
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png gAMA error crc.");
            return gAMA;
        }
        internal struct_png_pHYs read_pHYs(Stream sr, struct_png_datablock b)
        {
            var pHYs = new struct_png_pHYs { };

            pHYs.length = b.length;
            pHYs.tag = b.tag;
            pHYs.x_axis = helper.Readuint(sr);
            pHYs.y_axis = helper.Readuint(sr);
            pHYs.unit_specifier = helper.Readbyte(sr);
            if(helper.Readuint(sr) != b.crc) throw new ArgumentException("png pHYs error crc.");
            return pHYs;
        }
        internal struct_png_IDAT read_IDAT(Stream sr, struct_png_datablock b, Stream ot)
        {
            var IDAT = new struct_png_IDAT { };
            IDAT.length = b.length;
            IDAT.tag = b.tag;
            if (ot.Position > 0) // 只去除首个数据包的压缩参数
            {
                helper.ReadToMemory(sr, ot, IDAT.length);
                if (helper.Readuint(sr) != b.crc) throw new ArgumentException("png IDAT error crc.");
                return IDAT;
            }
            IDAT.zlib_compression = helper.Readbyte(sr);
            IDAT.additional_flags = helper.Readbyte(sr);
            helper.ReadToMemory(sr, ot, IDAT.length -2);
            if (helper.Readuint(sr) != b.crc) throw new ArgumentException("png IDAT error crc.");
            return IDAT;
        }

        public void Dispose() //释放内存是必要的
        {
            if (imgdata != null) imgdata.Dispose();
            imgdata = null;
        }
    }
}
